package com.test.common;

import java.math.BigInteger;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;

public class Memoizer<A, V> implements Computable<A, V> {
    private final Computable<A, V> computer;
    private final ConcurrentMap<A, FutureTask<V>> cache = new ConcurrentHashMap<>();
    private final ConcurrentMap<A, Long> computedTimes = new ConcurrentHashMap<>();
    private static final int EXPIRATION_TIME = 2000;

    public Memoizer(Computable<A, V> computer) {
        this.computer = computer;
    }

    @Override
    public V compute(final A arg) throws InterruptedException {
        while (true) {
            Long computedTime = computedTimes.get(arg);
            if (computedTime != null && computedTime + EXPIRATION_TIME < System.currentTimeMillis()) {
                System.out.println("removing");
                cache.remove(arg);
            }

            Future<V> f = cache.get(arg);
            if (f == null) {
                FutureTask<V> ft = new FutureTask<>(new Callable<V>() {
                    @Override
                    public V call() throws Exception {
                        System.out.println("computing");
                        V v = computer.compute(arg);
                        computedTimes.put(arg, System.currentTimeMillis());
                        return v;
                    }
                });

                f = cache.putIfAbsent(arg, ft);
                if (f == null) {
                    f = ft;
                    ft.run();
                }
            }

            try {
                return f.get();
            } catch (ExecutionException e) {
                cache.remove(arg);
                throw new RuntimeException("Computation failed", e);
            } catch (CancellationException e) {
                cache.remove(arg);
            }
        }
    }
}

public class ExpensiveComputation implements Computable<String, BigInteger> {

    @Override
    public BigInteger compute(String arg) throws InterruptedException {
        return new BigInteger(arg);
    }
}